RDBMSの実装を学ぶためにRust製OSSのtoyDBにTHANKYOU文とPLEASE句を実装してみた

時代に即したRDBMSの新機能？？ PLEASE句を実装してみました

#Rust

岩田智哉

2022.05.03

この記事は公開されてから1年以上経過しています。情報が古い可能性がありますので、ご注意ください。

MAD事業部@大阪の岩田です。

2021年の話になってしまいますが、エイプリルフールネタとして各種RDBMSにPLEASE句を実装する試みが流行していた？ようです。

この記事を読み、自分でもPLEASE句の実装に挑戦してRDBMSの実装について学んでみようと思い立ちました。Rustを勉強したいなーと思っていたこともありRust製の手頃なOSSを探してみたたところ、toyDBというちょうど良さそうなOSSを見つけたので、こちらを利用して進めていきます。

このブログでは最終的には以下を目標としてtoyDBに機能追加していきたいと思います。

THANKYOU文の実装
- クライアントがTHANKYOUとリクエストすると、Your Welcome!!を返却する
PLEASE句の実装
- SELECT文の先頭にPLEASEを付与できるようにする ※DELETEやUPDATE等は今回割愛
- PLEASE句を付与しないSELECT文は常にSeq Scanを選択するようにする

今回実装したコードは以下で公開しています。

https://github.com/cm-iwata/toydb/tree/feat/impl-thankyou-and-please

※Rust初心者なので良くない書き方をしてるかもしれませんがご了承下さい。

toyDBとは？

学習を目的としたRust製のOSS分散データベースでGitHubのリポジトリは以下です。

https://github.com/erikgrinaker/toydb/

toyDBには以下のような機能が実装されています。

Raftベースの分散合意アルゴリズム
MVCCをベースとしたACID準拠のトランザクションエンジン
B+ツリーとログ構造を備えたストレージエンジン
ヒューリスティックな最適化手法とタイムトラベルをサポートするクエリエンジン
JOINやGROUP BYを含む基本的なSQLインターフェース

公式ページに記載されているアーキテクチャ図はこちらです。

RDBMSの内部実装について学びたい人にとっては非常に良いサンプル実装ではないでしょうか？

環境

以下の環境で動作を確認しています

ベースにしたtoyDBのコミットハッシュ: 81cc7f24d7759d6a5883c1383793095f715015ec
OS: Mac OS X 10.15
rustc: 1.57.0

軽くtoyDBの動作確認

まずは軽くtoyDBを触ってみましょう。

GitHubからtoyDBをクローン

$ git clone https://github.com/cm-iwata/toydb.git

toyDBのサーバーを起動します。

$ (cd clusters/local && ./run.sh)

toydb-a 06:15:46 [INFO] Listening on 0.0.0.0:9601 (SQL) and 0.0.0.0:9701 (Raft)
toydb-a 06:15:46 [ERROR] Failed connecting to Raft peer 127.0.0.1:9703: Connection refused (os error 61)
toydb-a 06:15:46 [ERROR] Failed connecting to Raft peer 127.0.0.1:9705: Connection refused (os error 61)
toydb-a 06:15:46 [ERROR] Failed connecting to Raft peer 127.0.0.1:9704: Connection refused (os error 61)
toydb-a 06:15:46 [ERROR] Failed connecting to Raft peer 127.0.0.1:9702: Connection refused (os error 61)
toydb-b 06:15:46 [INFO] Replaying log entries 1 to 23
toydb-e 06:15:46 [INFO] Replaying log entries 1 to 23
toydb-d 06:15:46 [INFO] Replaying log entries 1 to 23
toydb-c 06:15:46 [INFO] Replaying log entries 1 to 23
toydb-e 06:15:46 [INFO] Listening on 0.0.0.0:9605 (SQL) and 0.0.0.0:9705 (Raft)
toydb-d 06:15:46 [INFO] Listening on 0.0.0.0:9604 (SQL) and 0.0.0.0:9704 (Raft)
toydb-e 06:15:46 [ERROR] Failed connecting to Raft peer 127.0.0.1:9703: Connection refused (os error 61)
toydb-b 06:15:46 [INFO] Listening on 0.0.0.0:9602 (SQL) and 0.0.0.0:9702 (Raft)
toydb-e 06:15:46 [ERROR] Failed connecting to Raft peer 127.0.0.1:9704: Connection refused (os error 61)
toydb-c 06:15:46 [INFO] Listening on 0.0.0.0:9603 (SQL) and 0.0.0.0:9703 (Raft)
toydb-d 06:15:47 [INFO] Starting election for term 31
toydb-e 06:15:47 [INFO] Discovered new term 31, following leader toydb-d
toydb-a 06:15:47 [INFO] Discovered new term 31, following leader toydb-d
toydb-c 06:15:47 [INFO] Discovered new term 31, following leader toydb-d
toydb-b 06:15:47 [INFO] Discovered new term 31, following leader toydb-d
toydb-e 06:15:47 [INFO] Voting for toydb-d in term 31 election
toydb-b 06:15:47 [INFO] Voting for toydb-d in term 31 election
toydb-c 06:15:47 [INFO] Voting for toydb-d in term 31 election
toydb-a 06:15:47 [INFO] Voting for toydb-d in term 31 election
toydb-d 06:15:47 [INFO] Won election for term 31, becoming leader

しばらく待つとリーダーの選出が完了し、...becoming leaderというログが出力され準備完了です。

クライアントツールであるtoysqlを起動し、toyDBサーバーに接続します。

$ cargo run --bin=toysql -- --host localhost --port 9601 -H

Connected to toyDB node "toydb-a". Enter !help for instructions.

ドキュメントに記載されているSQLを実行し、簡単にテーブル作成とデータのINSERTが実行できることを確認します。

toydb> CREATE TABLE movie (
    id INTEGER PRIMARY KEY,
    title STRING NOT NULL,
    release_year INTEGER INDEX,
    imdb_id STRING INDEX UNIQUE,
    bluray BOOLEAN NOT NULL DEFAULT TRUE
);
Created table movie

toydb> INSERT INTO movie
    (id, title, release_year)
VALUES
    (1, 'Sicario', 2015),
    (2, 'Stalker', 1979),
    (3, 'Her', 2013);
Created 3 rows

INSERTしたデータをSELECTしてみます

toydb> SELECT * FROM movie;
id|title|release_year|imdb_id|bluray
1|Sicario|2015|NULL|TRUE
2|Stalker|1979|NULL|TRUE
3|Her|2013|NULL|TRUE

これから実装するTHANKYOU文とPLEASE句がエラーになることを確認します

toydb> THANKYOU;
Error: Unexpected token thankyou

toydb> PLEASE SELECT * FROM movie;
Error: Unexpected token please

PLEASE句なしでも主キー、インデックスを利用した実行計画が選択されることを確認します

toydb> EXPLAIN SELECT * FROM movie WHERE id = 1;
KeyLookup: movie (1)

toydb> EXPLAIN SELECT * FROM movie WHERE release_year = 2015;
IndexLookup: movie column release_year (2015)

これでウォーミングアップ完了です。

THANKYOU文を実装してみる

まずは簡単そうなところでTHANKYOU文を実装していきます。

前述のアーキテクチャ図を見る限りSQL Engineに手を加えると良さそうです。ディレクトリ構成を見る限りsrc/sql配下にSQL Engine周りのコードが置かれてそうです。

$ ls -l src/sql/
total 32
drwxr-xr-x  5 iwata.tomoya  staff    160 May  3 15:33 engine
drwxr-xr-x  9 iwata.tomoya  staff    288 May  3 15:34 execution
-rw-r--r--  1 iwata.tomoya  staff     96 May  1 23:04 mod.rs
drwxr-xr-x  5 iwata.tomoya  staff    160 May  3 15:34 parser
drwxr-xr-x  5 iwata.tomoya  staff    160 May  3 15:34 plan
-rw-r--r--  1 iwata.tomoya  staff  10464 May  1 23:04 schema.rs
drwxr-xr-x  4 iwata.tomoya  staff    128 May  1 23:04 types

以後は変更箇所の一部を取り上げながら紹介していきます。変更箇所の全体像は以下のコミットを参照して下さい。

https://github.com/cm-iwata/toydb/commit/1abeb22118426d8a077e04d3064181b9baa1a903

まずはParserの修正からです。Keywordというenumに予約語が定義されているので、ここにThankyouを追加します。

/// Lexer keywords
#[derive(Clone, Debug, PartialEq)]
pub enum Keyword {
    And,
    As,
    Asc,
    ...略
    Thankyou,
    ...略

StatementというenumにSQL文が定義されているので、ここにもThankyouを追加します。

/// Statements
#[derive(Clone, Debug, PartialEq)]
#[allow(clippy::large_enum_variant)]
pub enum Statement {
    Begin {
        readonly: bool,
        version: Option<u64>,
    },
    ...略
    Thankyou,
}

この調子でパース処理を実装し、最後にSQL文の実行処理にTHANKYOU文の実行処理を追加します。ここではTHANKYOU文には固定でYour Welcome!!を返却する実装としました。

impl<E: Engine + 'static> Session<E> {
    /// Executes a query, managing transaction status for the session
    pub fn execute(&mut self, query: &str) -> Result<ResultSet> {
        // FIXME We should match on self.txn as well, but get this error:
        // error[E0009]: cannot bind by-move and by-ref in the same pattern
        // ...which seems like an arbitrary compiler limitation
        match Parser::new(query).parse()? {
            ast::Statement::Thankyou => {
                Ok(ResultSet::Thankyou { message: "Your Welcome!!".to_string() })
            }

テスト

実装できたらテストしてみましょう。先程と同様にクライアントツールのtoysqlでtodDBサーバーに接続し、THANKYOU文を実行します。

toydb> thankyou;
Your Welcome!!

Your Welcome!!が返却されました。THANKYOU文の実装成功です。

PLEASE句を実装してみる

続いてPLEASE句の実装です。

SELECT文の先頭にPLEASE句を付与できるように修正

まずはSELECT文の先頭にPLEASE句を付与してもエラーにならないようにします。該当のコミットはこちらです。

https://github.com/cm-iwata/toydb/commit/63f8a97202d50d355cb2d7497ad2cf631eded8fb

まずは先程のTHANKYOU文と同様に予約後にPLEASEを追加します。

/// Lexer keywords
#[derive(Clone, Debug, PartialEq)]
pub enum Keyword {
    And,
    As,
    Asc,
    ...略
    Please,

SELECT文のASTにPLEASE句の有無を追加します。

/// Statements
#[derive(Clone, Debug, PartialEq)]
#[allow(clippy::large_enum_variant)]
pub enum Statement {
		...略
    Select {
        with_please: bool,
        select: Vec<(Expression, Option<String>)>,
        from: Vec<FromItem>,
        r#where: Option<Expression>,
        group_by: Vec<Expression>,
        having: Option<Expression>,
        order: Vec<(Expression, Order)>,
        offset: Option<Expression>,
        limit: Option<Expression>,
    },

PLEASE句の指定有無だけ分かれば良いので、with_please: boolの定義を追加しています。

続いてSQLのパース処理を修正します。まずはPLEASE句のパース処理を追加

PLEASE句が指定された場合はPLEASE句が指定されたことを表すbool値にtrueを設定し、後続のSQLを再帰的にパースしていきます。

    fn parse_statement_please(&mut self) -> Result<ast::Statement> {
        self.next_expect(Some(Keyword::Please.into()))?;

        match self.peek()? {
            Some(Token::Keyword(Keyword::Select)) => self.parse_statement(true),
            Some(token) => Err(Error::Parse(format!("Unexpected token {}", token))),
            None => Err(Error::Parse("Unexpected end of input".into())),
        }
    }

今回はSELECT文のみPLEASE句の付与を許可するため、

SELECT文
EXPLAIN文

のパース処理も合わせて修正しています。

    /// Parses an SQL statement
    fn parse_statement(&mut self, with_please: bool) -> Result<ast::Statement> {
        match self.peek()? {
            Some(Token::Keyword(Keyword::Please)) => self.parse_statement_please(),
						...略
            Some(Token::Keyword(Keyword::Select)) => self.parse_statement_select(with_please),

これでPLEASE句付きのSELECT文が実行できるようになりました。

PLEASE句無しのSELECT文は常にSeq Scanを利用するように修正

続いてPLEASE句無しのSELECT文は常にSeq Scanを利用するように修正します。

該当のコミットはこちらです。

https://github.com/cm-iwata/toydb/commit/c1aa308b39561c548b10fdcb7b234c8208652cb9

まずは実行計画のノードScanにPLEASE句の有無を表すbool値を追加します。このbool値はSQLをパースした際のPLEASE有無に応じて値を設定します。

/// A plan node
#[derive(Debug, PartialEq, Serialize, Deserialize)]
pub enum Node {
    ...略
    Scan {
        table: String,
        alias: Option<String>,
        filter: Option<Expression>,
        with_please: bool,
    },

続いて最適化処理です。元々の挙動では、インデックスが利用可能な場合は実行計画をSeq ScanからIndex Lookupに変更してくれるのですが、PLEASE句の指定がない場合はインデックスの利用可否をいちいち判定せずに、そのままSeq Scanの実行計画を返却するようにします。

impl<'a, C: Catalog> Optimizer for IndexLookup<'a, C> {
    fn optimize(&self, node: Node) -> Result<Node> {
        node.transform(&|n| Ok(n), &|n| match n {
            Node::Scan { table, alias, filter: Some(filter), with_please } => {
                if !with_please {
                    return Ok(Node::Scan { table, alias, filter: Some(filter), with_please });
                }

これでPLEASE句なしの失礼な依頼に対しては常にSeq Scanしか実行されないようになりました。

EXPLAIN結果にPLEASE句の有無を出力するよう修正

最後にEXPLAIN結果からPLEASE句の指定有無が表示されるように修正を加えておきます。これでPLEASE句の指定を忘れによって意図せずSeq Scanが選択されているクエリもデバッグするのが簡単になります。

該当コミットはこちらです。

https://github.com/cm-iwata/toydb/commit/c9c7b7d6b62e775f4dcad9a85a78c3c2bff9902e

    // Displays the node, where prefix gives the node prefix.
    pub fn format(&self, mut indent: String, root: bool, last: bool) -> String {
        let mut s = indent.clone();
        if !last {
            s += "├─ ";
            indent += "│  "
        } else if !root {
            s += "└─ ";
            indent += "   ";
        }
        match self {
            ...略
            Self::Scan { table, alias, filter, with_please } => {
                s += &format!("Scan: {} with please: {}", table, with_please);
                if let Some(alias) = alias {
                    s += &format!(" as {}", alias);
                }
                if let Some(expr) = filter {
                    s += &format!(" ({})", expr);
                }
                s += "\n";
            }

テスト

toysqlからテストしてみます。まずソースコード修正前と同様に主キー/インデックスが設定されたカラムで絞り込みを行うSELECT文をEXPLAINしてみます。

toydb> EXPLAIN SELECT * FROM movie WHERE id = 1;
Scan: movie with please: false (id = 1)

toydb> EXPLAIN SELECT * FROM movie WHERE release_year = 2015;
Scan: movie with please: false (release_year = 2015)

実行計画がScanに変わっていることが分かります。続いてPLEASE句を付与してみましょう。

toydb> EXPLAIN PLEASE SELECT * FROM movie WHERE id = 1;
KeyLookup: movie (1)
toydb>  EXPLAIN PLEASE SELECT * FROM movie WHERE release_year = 2015;
IndexLookup: movie column release_year (2015)

今度はちゃんと主キー/インデックスを利用する実行計画に変わってくれました。toysqlはテーブルの中身をメモリに乗せるので、Seq Scanになっても処理速度の低下が分かりにくいですが、10万レコード投入してからPLEASE句あり/無しそれぞれのSQLを実行すると、PLEASE句無しの場合は体感で「ん？ちょっと遅いかな？」ぐらいの変化が見られました。

※ちゃんと計測するのは少し面倒だったので今回は省略します。

ちなみに... cargo testを実行するとオプティマイザやプランナ関連のテストが大量に失敗します。テストケースにはPLEASE句が指定されていないため仕方ないですね...面倒なので今回はテストケースの修正は省略しています。

まとめ

Rust初心者ですが、楽しみながら実装できました。今回はTHANKYOU文とPLEASE句という業務には役立たない機能追加に挑戦しましたが、それでも普段RDBMSが実行している処理の裏側の理解を進めることができたのは収穫です。toyDBはRDBMSの内部実装を学習するために非常に良い題材だなという印象を受けたので、RDBMSについて掘り下げて学んでみたい人はぜひtoyDBで遊んで見て下さい。